#!/bin/bash

# FLUX训练环境设置脚本
# 一键准备FLUX训练所需的所有文件和目录

echo "========================================"
echo "FLUX训练环境设置脚本"
echo "========================================"

# 设置基础路径
WORKSPACE="/workspace"
TRAINING_DIR="$WORKSPACE/training_data_all"

echo "设置训练环境..."
echo "工作空间: $WORKSPACE"
echo "训练目录: $TRAINING_DIR"

# 1. 创建目录结构
echo
echo "1. 创建目录结构..."
mkdir -p "$TRAINING_DIR/images"
mkdir -p "$TRAINING_DIR/output"
mkdir -p "$TRAINING_DIR/reg_images"

echo "✓ 创建目录: $TRAINING_DIR/images"
echo "✓ 创建目录: $TRAINING_DIR/output"
echo "✓ 创建目录: $TRAINING_DIR/reg_images"

# 2. 复制数据集配置文件
echo
echo "2. 复制数据集配置文件..."
if [ -f "dataset_1024_bs2.toml" ]; then
    cp "dataset_1024_bs2.toml" "$TRAINING_DIR/"
    echo "✓ 复制: dataset_1024_bs2.toml"
else
    echo "⚠️  文件不存在: dataset_1024_bs2.toml"
fi

if [ -f "dataset_1024_bs1.toml" ]; then
    cp "dataset_1024_bs1.toml" "$TRAINING_DIR/"
    echo "✓ 复制: dataset_1024_bs1.toml"
else
    echo "⚠️  文件不存在: dataset_1024_bs1.toml"
fi

# 3. 复制采样提示词文件
echo
echo "3. 复制采样提示词文件..."
if [ -f "sample_prompts.txt" ]; then
    cp "sample_prompts.txt" "$TRAINING_DIR/"
    echo "✓ 复制: sample_prompts.txt"
else
    echo "⚠️  文件不存在: sample_prompts.txt"
fi

# 4. 创建示例标题文件
echo
echo "4. 创建示例标题文件..."
cat > "$TRAINING_DIR/example_caption.txt" << 'EOF'
# FLUX训练标题文件示例
# 每个图像都需要一个对应的.txt标题文件，文件名与图像相同（除了扩展名）

# 示例标题格式：
masterpiece, best quality, (1girl), in white shirts, upper body, looking at viewer, simple background --n low quality, worst quality, bad anatomy,bad composition, poor, low effort

# 标题文件命名规则：
# 图像文件: image001.jpg
# 标题文件: image001.txt

# 标题内容建议：
# - 使用英文描述
# - 包含图像的主要特征
# - 使用正向提示词和负向提示词
# - 保持描述准确和一致
EOF

echo "✓ 创建示例标题文件: $TRAINING_DIR/example_caption.txt"

# 5. 创建README文件
echo
echo "5. 创建README文件..."
cat > "$TRAINING_DIR/README.md" << 'EOF'
# FLUX训练数据集说明

## 目录结构
```
training_data_all/
├── images/              # 训练图像目录
├── output/              # 训练输出目录
├── reg_images/          # 正则化图像目录（可选）
├── dataset_1024_bs2.toml    # LoRA训练配置
├── dataset_1024_bs1.toml    # 微调训练配置
├── sample_prompts.txt       # 采样提示词
└── example_caption.txt      # 标题文件示例
```

## 使用步骤

### 1. 准备训练图像
- 将训练图像放入 `images/` 目录
- 支持的格式：jpg, jpeg, png, webp, bmp
- 建议分辨率：512x512 到 2048x2048

### 2. 创建标题文件
- 为每个图像创建对应的.txt标题文件
- 文件名与图像相同（除了扩展名）
- 参考 `example_caption.txt` 的格式

### 3. 运行训练脚本
```bash
# LoRA训练
./flux_24g_training.sh

# 高级LoRA训练
./flux_24g_optimized.sh

# 微调训练
./flux_24g_finetune.sh
```

## 标题文件格式
```
masterpiece, best quality, (1girl), in white shirts, upper body, looking at viewer, simple background --n low quality, worst quality, bad anatomy,bad composition, poor, low effort
```

## 注意事项
- 确保所有图像都有对应的标题文件
- 标题描述要准确和一致
- 使用英文描述
- 包含正向和负向提示词
EOF

echo "✓ 创建README文件: $TRAINING_DIR/README.md"

# 6. 检查现有文件
echo
echo "6. 检查现有文件..."
if [ -d "$TRAINING_DIR/images" ]; then
    IMAGE_COUNT=$(find "$TRAINING_DIR/images" -type f \( -name "*.jpg" -o -name "*.jpeg" -o -name "*.png" -o -name "*.webp" -o -name "*.bmp" \) | wc -l)
    echo "✓ 图像目录存在，包含 $IMAGE_COUNT 个图像文件"
else
    echo "⚠️  图像目录不存在"
fi

# 7. 设置权限
echo
echo "7. 设置文件权限..."
chmod +x flux_24g_training.sh 2>/dev/null || echo "⚠️  flux_24g_training.sh 不存在"
chmod +x flux_24g_optimized.sh 2>/dev/null || echo "⚠️  flux_24g_optimized.sh 不存在"
chmod +x flux_24g_finetune.sh 2>/dev/null || echo "⚠️  flux_24g_finetune.sh 不存在"

echo "✓ 设置脚本执行权限"

# 8. 显示完成信息
echo
echo "========================================"
echo "训练环境设置完成！"
echo "========================================"
echo
echo "下一步操作："
echo "1. 将训练图像放入: $TRAINING_DIR/images/"
echo "2. 为每个图像创建对应的.txt标题文件"
echo "3. 编辑标题文件，提供准确的图像描述"
echo "4. 运行训练脚本："
echo "   - LoRA训练: ./flux_24g_training.sh"
echo "   - 高级LoRA: ./flux_24g_optimized.sh"
echo "   - 微调训练: ./flux_24g_finetune.sh"
echo
echo "详细说明请查看: $TRAINING_DIR/README.md"
echo
echo "当前目录结构："
tree "$TRAINING_DIR" 2>/dev/null || ls -la "$TRAINING_DIR"
